import torch

input = torch.tensor([
    [1, 2, 0, 3, 1],
    [0, 1, 2, 3, 1],
    [1, 2, 1, 0, 0],
    [5, 2, 3, 1, 1],
    [2, 1, 0, 1, 1]
])

kernel = torch.tensor([
    [1, 2, 1],
    [0, 1, 0],
    [2, 1, 0]
])

print(input.shape)
print(kernel.shape)

input = torch.reshape(input, (1, 1, 5, 5))
kernel = torch.reshape(kernel, (1, 1, 3, 3))

output = torch.nn.functional.conv2d(input=input, weight=kernel, stride=1)
print(output)

output = torch.nn.functional.conv2d(input=input, weight=kernel, stride=2)
print(output)

output = torch.nn.functional.conv2d(input=input, weight=kernel, stride=1, padding=1)
print(output)
